package export;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import neuron.Coord;
import neuron.Dendrite;
import neuron.Neuron;
import neuron.Segment;
import neuron.Statistics;

import matlab.DoubleArrayVariable;
import matlab.Matlab5File;
import common.statistics.Series;

public class NeuronToMatlabExport {

	static class PointAdder
	{
		List<Coord> points = new ArrayList<Coord>();
		int idx = 0;
		public int addPoint(Coord c)
		{
			points.add(c);
			return idx++;
		}
	}
	
	static class Edge 
	{
		int a, b;
		public Edge(int a, int b) {
			this.a = a;
			this.b = b;
		}
	}
	
	File file;
	Neuron neur;
	double[] planes;
	PointAdder points = new PointAdder();
	
	List<Integer> basalRoots = new ArrayList<Integer>();
	int apicalRoot = -1;		
	
	List<Integer> branchPoints = new ArrayList<Integer>();
	List<Integer> terminalPoints = new ArrayList<Integer>();
	List<Edge> edges = new ArrayList<Edge>();
	
	
	public NeuronToMatlabExport(File file, Neuron neur, double[] planes) {
		this.file = file; this.neur = neur; this.planes = planes;
	}
	 
	/**
	 * Export segment. 
	 * 
	 * @param s
	 * @param prevCoordIdx index of previous coordinate, or -1 for root segment
	 */
	private void exportSegment(Segment s, int prevCoordIdx) 
	{
		exportSegment(s, prevCoordIdx, false);
	}

	/**
	 * Export segment
	 * @param s
	 * @param prevCoordIdx
	 * @param isApical
	 */
	private void exportSegment(Segment s, int prevCoordIdx, boolean isApical)
	{
		// Add root ?
		if (prevCoordIdx == -1) {
			prevCoordIdx = points.addPoint(s.firstCoordinate());
			if (!isApical) basalRoots.add(prevCoordIdx);
			else apicalRoot = prevCoordIdx;
		}
		
		// add this segment
		int thisIdx = points.addPoint(s.lastCoordinate());
		if (s.isTerminalSegment()) {
			terminalPoints.add(thisIdx);
		} else {
			branchPoints.add(thisIdx); // TODO FIXME check if there is also continuations
		}
		edges.add(new Edge(prevCoordIdx, thisIdx));
		
		// recurse to branches
		if (s.getBranch(0) != null) { exportSegment(s.getBranch(0), thisIdx); }
		if (s.getBranch(1) != null) { exportSegment(s.getBranch(1), thisIdx); }
		if (s.getBranch(2) != null) { exportSegment(s.getBranch(2), thisIdx); }
	}
	
	public void export() throws IOException
	{
		Matlab5File mat = new Matlab5File();
		
		// soma
		mat.addDoubleArray("soma_coordinate", new double[] { 0, 0, 0 }, 1, 3);
		
		// planes
		if (planes != null && planes.length >= 4) {
			mat.addDouble("soma_pia", planes[0]);
			mat.addDouble("soma_l4l5", planes[1]);
			mat.addDouble("soma_wm", planes[3]);
		}

		// source file
		if (neur.getMetadata("FILE") != null)
			mat.addText("neuron", (String) neur.getMetadata("FILE"));

		
		// COLLECT DATA
		
		// go through all dendrites
		for (Dendrite d : neur.basalDendrites()) {
			exportSegment(d.getRootSegment(), -1);
		}

		// also apical
		if (neur.getApicalDendrite() != null) {
			exportSegment(neur.getApicalDendrite().getRootSegment(), -1, true);
		}
		
		
		// WRITE DATA
		
		// add points
		int m = points.points.size();
		double[] pointArray = new double[3*m];
		for (int i = 0; i < m; i++) {
			Coord c  = points.points.get(i);
			pointArray[i*3 + 0] = c.x;
			pointArray[i*3 + 1] = c.y;
			pointArray[i*3 + 2] = c.z;
		}
		mat.addDoubleArray("coordinates", DoubleArrayVariable.transpose(pointArray, 3, m), m, 3); 
		
		// add root indices
		double[] rootArr = new double[basalRoots.size()];
		for (int i = 0; i < basalRoots.size(); i++) rootArr[i] = basalRoots.get(i);
		mat.addDoubleArray("basal_roots", rootArr, rootArr.length, 1); 
		
		mat.addDouble("apical_root", apicalRoot);
		
		// add branch points
		double[] branchArr = new double[branchPoints.size()];
		for (int i = 0; i < branchPoints.size(); i++) branchArr[i] = branchPoints.get(i);
		mat.addDoubleArray("branch_points", branchArr, branchArr.length, 1); 

		// add terminal points
		double[] termArr = new double[terminalPoints.size()];
		for (int i = 0; i < terminalPoints.size(); i++) termArr[i] = terminalPoints.get(i);
		mat.addDoubleArray("terminal_points", termArr, termArr.length, 1); 

		// edges
		double[] edgeArr = new double[edges.size()*2];
		for (int i = 0; i < edges.size(); i++) {
			edgeArr[i*2+0] = edges.get(i).a; 
			edgeArr[i*2+1] = edges.get(i).b; 
		}
		mat.addDoubleArray("edges", DoubleArrayVariable.transpose(edgeArr, 2, edges.size()), edges.size(), 2); 
		
		// TODO add basal_termseg_lengths, basal_interseg_lengths, oblique_, tuft_ 
		Series basalTermSegLens = new Series();
		Series basalIntSegLens = new Series();
		for (Dendrite d : neur.basalDendrites()) {
			Statistics stat = new Statistics(d.root);
			basalTermSegLens.merge(stat.terminalSegmentLengths());
			basalIntSegLens.merge(stat.intermediateSegmentLengths());
		}
		mat.addDoubleArray("basal_intseg_lengths", basalIntSegLens.values(), basalIntSegLens.count(), 1);
		mat.addDoubleArray("basal_termseg_lengths", basalTermSegLens.values(), basalTermSegLens.count(), 1);
		
		
		// write to file
		mat.write(file);
	}
}
